import os
import argparse

# hosts 文件的路径
hosts_path = r"C:\Windows\System32\drivers\etc\hosts"

def update_hosts(ip, host):
    try:
        # 读取当前的 hosts 文件内容
        with open(hosts_path, 'r', encoding='utf-8') as file:
            lines = file.readlines()

        # 移除现有的指定域名相关条目
        new_lines = []
        for line in lines:
            if host not in line:
                new_lines.append(line)

        # 添加新的指定 IP 和域名映射条目
        new_entry = f"{ip} {host}\n"
        new_lines.append(new_entry)

        # 将更新后的内容写回 hosts 文件
        with open(hosts_path, 'w', encoding='utf-8') as file:
            file.writelines(new_lines)

        print("hosts 文件更新成功！")
    except Exception as e:
        print(f"更新 hosts 文件时出错: {e}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description='此脚本用于更新 Windows 系统的 hosts 文件，添加或修改指定的 IP 地址与域名映射。',
        epilog='示例：python addhosts.py --ip 1.2.3.4 --host example.com'
    )
    parser.add_argument('--ip', required=True, help='要添加到 hosts 文件中的 IP 地址')
    parser.add_argument('--host', required=True, help='要添加到 hosts 文件中的域名')
    args = parser.parse_args()

    update_hosts(args.ip, args.host)